import os
import cv2
import numpy as np
import streamlit as st
import ShapeRecognition as Sr
import ImageCorrection as Ic
import LaneLineDetection as Lld


# 初始化会话状态：保存跨刷新的关键数据（避免参数/图像丢失）
def init_session_state():
    # 车道线掩码参数（默认值与原代码一致）
    if "lane_params" not in st.session_state:
        st.session_state.lane_params = {
            "first_width": 0.0, "first_height": 0.5,
            "second_width": 0.5, "second_height": 0.33,
            "third_width": 0.5, "third_height": 0.33,
            "fourth_width": 1.0, "fourth_height": 0.5
        }
    # 上次上传文件路径（避免重复读取）
    if "last_file_path" not in st.session_state:
        st.session_state.last_file_path = ""
    # 车道线预处理数据（缩放后图像、边缘图、宽高）
    if "lane_frame_data" not in st.session_state:
        st.session_state.lane_frame_data = {
            "edges_img": None, "height": 0, "width": 0, "raw_frame": None
        }
    # 当前用户输入（文本+文件，刷新后保留显示）
    if "current_user_input" not in st.session_state:
        st.session_state.current_user_input = None


# 初始化会话状态（确保页面加载时变量已定义）
init_session_state()

# 处理用户聊天输入（仅在提交新内容时执行，避免刷新刷新重复触发）
if user_input := st.chat_input("来和我聊天吧", accept_file=True, file_type=['jpg', 'png', 'jpeg']):
    # 保存当前用户输入到会话状态（刷新后可复用）
    st.session_state.current_user_input = user_input

    # 处理文件上传与保存（仅新文件时执行）
    if user_input["files"]:
        # 确保保存目录存在
        if not os.path.exists("./model_save_img"):
            os.mkdir("./model_save_img")
        # 构建文件路径
        file_path = os.path.join("./model_save_img", user_input["files"][0].name)
        # 读取并保存图像（保持原代码逻辑）
        file_bytes = user_input["files"][0].read()
        np_arr = np.frombuffer(file_bytes, np.uint8)
        img = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
        cv2.imwrite(file_path, img)

        # 更新会话状态：保存文件路径和车道线预处理结果
        st.session_state.last_file_path = file_path
        # 提前预处理车道线图像（避免参数变化时重复计算）
        frame = cv2.imread(file_path)
        frame_resized = cv2.resize(frame, (0, 0), fx=0.3, fy=0.3)  # 保持原缩放比例
        edges_img, height, width = Lld.img_preprocess(frame_resized)
        st.session_state.lane_frame_data = {
            "edges_img": edges_img, "height": height,
            "width": width, "raw_frame": frame_resized
        }

# 显示用户消息（从会话状态读取，刷新后不消失）
if st.session_state.current_user_input:
    user_input = st.session_state.current_user_input
    with st.chat_message("user"):
        # 显示用户输入文本
        if user_input.text:
            st.markdown(user_input.text)
        # 显示上传的原始图像
        if user_input["files"]:
            st.image(user_input["files"][0], width=300)

# 显示助手回复（核心逻辑：从会话状态读数据，避免刷新空白）
if st.session_state.current_user_input and st.session_state.last_file_path:
    user_input = st.session_state.current_user_input
    file_path = st.session_state.last_file_path  # 从会话状态取文件路径
    with st.chat_message("assistant"):
        # 1. 灰度处理（保留原逻辑）
        if "灰度" in user_input.text:
            img = cv2.imread(file_path)
            img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            st.markdown("灰度处理结果如下：")
            st.image(img_gray, width=300)

        # 2. 浮雕/滤镜处理（优化边界判断为np.clip，更简洁）
        if "浮雕" in user_input.text or "滤镜" in user_input.text:
            img = cv2.imread(file_path)
            h, w, _ = img.shape
            img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            img_FuDiao = np.zeros_like(img)
            for i in range(h):
                for j in range(w - 2):
                    p0 = img_gray[i, j]
                    p2 = img_gray[i, j + 2]
                    newP = p0 - p2 + 150
                    newP = np.clip(newP, 0, 255)  # 替代if-elif，避免越界
                    img_FuDiao[i, j] = newP
            st.markdown("浮雕处理结果如下：")
            st.image(img_FuDiao, width=300)

        # 3. 二值处理（保留原逻辑）
        if "二值" in user_input.text:
            img = cv2.imread(file_path)
            img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            img_binary = cv2.threshold(img_gray, 150, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)[1]
            st.markdown("二值处理结果如下：")
            st.image(img_binary, width=300)

        # 4. 旋转处理（保留原逻辑）
        if "旋转" in user_input.text:
            img = cv2.imread(file_path)
            h, w, _ = img.shape
            M = cv2.getRotationMatrix2D((w / 2, h / 2), 45, 1)
            img_rotate = cv2.warpAffine(
                img, M, (w, h), borderMode=cv2.BORDER_CONSTANT, borderValue=(255, 255, 255)
            )
            st.markdown("旋转处理结果如下：")
            st.image(cv2.cvtColor(img_rotate, cv2.COLOR_BGR2RGB), width=300)

        # 5. 边缘处理（保留原逻辑）
        if "边缘" in user_input.text:
            img = cv2.imread(file_path)
            img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            img_binary = cv2.threshold(img_gray, 150, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)[1]
            img_gauss = cv2.GaussianBlur(img_binary, (5, 5), 0)
            img_edge = cv2.Canny(img_gauss, 50, 110)
            st.markdown("边缘处理结果如下：")
            st.image(img_edge, width=300)

        # 6. 形状识别（保留原逻辑）
        if "形状" in user_input.text or "识别" in user_input.text:
            img = Sr.shape_recognition(file_path)
            st.markdown("形状识别结果如下：")
            st.image(cv2.cvtColor(img, cv2.COLOR_BGR2RGB), width=600)

        # 7. 图片矫正（保留原逻辑）
        if "矫正" in user_input.text:
            img = Ic.image_correction(file_path)
            st.markdown("图片矫正结果如下：")
            st.image(cv2.cvtColor(img, cv2.COLOR_BGR2RGB), width=400)

        # 8. 车道线识别（重点修复：参数从会话状态读取，刷新不丢失）
        if "车道" in user_input.text:
            lane_data = st.session_state.lane_frame_data  # 从会话状态取预处理数据
            # 防护：避免未上传图像时报错
            if lane_data["edges_img"] is None:
                st.warning("请先上传车道图像，再调整掩码参数")
                st.stop()

            # 显示原始图像（与原代码一致）
            st.image(user_input["files"][0], width=400)
            st.markdown("请输入掩码参数：")

            # 分两列布局参数输入框（优化视觉体验，减少滚动）
            col1, col2 = st.columns(2)
            with col1:
                # 第一个点：值从会话状态读取，修改后实时保存
                st.session_state.lane_params["first_width"] = st.number_input(
                    "第一个点位置--width * ？", min_value=0.0, max_value=1.0,
                    value=st.session_state.lane_params["first_width"], step=0.1
                )
                st.session_state.lane_params["first_height"] = st.number_input(
                    "第一个点位置--height * ？", min_value=0.0, max_value=1.0,
                    value=st.session_state.lane_params["first_height"], step=0.1
                )
                # 第二个点
                st.session_state.lane_params["second_width"] = st.number_input(
                    "第二个点位置--width * ？", min_value=0.0, max_value=1.0,
                    value=st.session_state.lane_params["second_width"], step=0.1
                )
                st.session_state.lane_params["second_height"] = st.number_input(
                    "第二个点位置--height * ？", min_value=0.0, max_value=1.0,
                    value=st.session_state.lane_params["second_height"], step=0.1
                )

            with col2:
                # 第三个点
                st.session_state.lane_params["third_width"] = st.number_input(
                    "第三个点位置--width * ？", min_value=0.0, max_value=1.0,
                    value=st.session_state.lane_params["third_width"], step=0.1
                )
                st.session_state.lane_params["third_height"] = st.number_input(
                    "第三个点位置--height * ？", min_value=0.0, max_value=1.0,
                    value=st.session_state.lane_params["third_height"], step=0.1
                )
                # 第四个点
                st.session_state.lane_params["fourth_width"] = st.number_input(
                    "第四个点位置--width * ？", min_value=0.0, max_value=1.0,
                    value=st.session_state.lane_params["fourth_width"], step=0.1
                )
                st.session_state.lane_params["fourth_height"] = st.number_input(
                    "第四个点位置--height * ？", min_value=0.0, max_value=1.0,
                    value=st.session_state.lane_params["fourth_height"], step=0.1
                )

            # 开始识别：用独立按钮触发（替代文本关键词，避免参数变化误触发）
            if st.button("开始车道线识别", type="primary"):
                # 从会话话状态取参数，简化调用
                p = st.session_state.lane_params
                h, w = lane_data["height"], lane_data["width"]
                # 计算ROI（移除原代码固定偏移，避免坐标负数；确保整数类型）
                ROI = np.array([
                    [(int(p["first_width"] * w), int(p["first_height"] * h)),
                     (int(p["second_width"] * w) - 20, int(p["second_height"] * h)),
                     (int(p["third_width"] * w) + 20, int(p["third_height"] * h)),
                     (int(p["fourth_width"] * w), int(p["fourth_height"] * h))]
                ], dtype=np.int32)

                # 执行车道线检测（保留原逻辑）
                masked_edges = Lld.img_mask(lane_data["edges_img"], ROI)
                img_result = Lld.line_detection(masked_edges)
                # 显示结果（转换为RGB格式，避免颜色偏色）
                st.markdown("车道线识别结果如下：")
                st.image(cv2.cvtColor(img_result, cv2.COLOR_BGR2RGB), width=600)
